# Function: Parse map file with GCC BUILD
# Coding:   utf-8
# Python version: v3.6.4
##############################################################################################################################

import os
import struct
import json
import re
import shutil
import time
import argparse

LEVEL_ERROR = '[Error]'
LEVEL_INFO = '[Info]'
LEVEL_WARN = '[Warning]'

class _Log(object):
  def __init__(self) -> None:
    pass

  def LOG(self, level, msg):
    def time_str():
      return time.strftime('%Y-%m-%d %H:%M:%S  ', time.localtime())
    log_msg = time_str() + level + ': ' + str(msg)
    print(log_msg)
    if level == LEVEL_ERROR:
      exit('Error: {}'.format(log_msg))


def _WorkspaceInit(output_path=None, log_printer=None) -> str:
  if output_path is None:
    output_path = './out' + '_' + time.strftime('%Y%m%d_%H%M%S', time.localtime()) + '/'
  else:
    output_path = output_path.rstrip('/').rstrip('\\') + '/'

  if os.path.isdir(output_path) == True:
    shutil.rmtree(output_path)

  try:
    os.mkdir(output_path)
  except Exception as e:
    log_printer(LEVEL_ERROR, 'Cannot make dir [{}], error=[{}]'.format(output_path, str(e)))

  return output_path


class BinEditor(object):
  def __init__(self, bin_path, map_path, var, mode, changes_path, LOG) -> None:
    self.bin_path = bin_path
    self.map_path = map_path
    self.variable_str = var
    self.workmode = mode
    self.variable_new_data_path = changes_path
    self.LOG = LOG
    self.SUPPORTED_MODE = ('r', 'w')
    self.map_lines = []
    self.bin_bytes = ''
    self.variable_addr = '0x0'
    self.variable_size = 0
    self.bin_start_addr = '0x0'
    self.variable_new_data_json = ''
    self.variable_bytes_start = 0
    self.variable_bytes_stop = 0
    self.CheckParamValid()

  def CheckParamValid(self):
    if not os.path.exists(self.bin_path):
      self.LOG(LEVEL_ERROR, 'The bin path not exist, please check: [{}]'.format(self.bin_path))

    if not os.path.exists(self.map_path):
      self.LOG(LEVEL_ERROR, 'The map path not exist, please check: [{}]'.format(self.bin_path))

    if not self.variable_str:
      self.LOG(LEVEL_ERROR, 'Please specify a right variable string to read/ write [{}]'.format(self.variable_str))

    if self.workmode not in self.SUPPORTED_MODE:
      self.LOG(LEVEL_ERROR, 'Please specify the mode you want to operate on the variable, supported mode is  [{}], but you set [{}]'.format(self.SUPPORTED_MODE, self.variable_str))

    if self.workmode == 'w':
      if not os.path.exists(self.variable_new_data_path):
        self.LOG(LEVEL_ERROR, 'Please specify the json file path you want to change on the variable. [{}]'.format(self.variable_new_path))


  def ReadFiles(self):
    try:
      fd = open(self.map_path, 'r', encoding='utf8')
      self.map_lines = fd.read().split('\n')
      fd.close()
    except Exception as e:
      self.LOG(LEVEL_ERROR, 'Can not open file: [{}], error=[{}]'.format(self.map_path, str(e)))

    try:
      fd = open(self.bin_path, 'rb')
      self.bin_bytes = fd.read()
      fd.close()
    except Exception as e:
      self.LOG(LEVEL_ERROR, 'Can not open file: [{}], error=[{}]'.format(self.bin_path, str(e)))

    if self.workmode == 'w':
      try:
        fd = open(self.variable_new_data_path, 'r', encoding='utf8')
        self.variable_new_data_json = fd.read()
      except Exception as e:
        self.LOG(LEVEL_ERROR, 'Can not open file: [{}], error=[{}]'.format(self.variable_new_data_path, str(e)))
    self.LOG(LEVEL_INFO, 'map file: [{}]'.format(self.map_path))
    self.LOG(LEVEL_INFO, 'app bin: [{}]'.format(self.bin_path))


  def FindVariableAddr(self):
    VAR_RE_PATTERN = r'^ \.\w+\.{}'.format(self.variable_str)
    VAR_RE_PATTERN_COMPILED = re.compile(VAR_RE_PATTERN)
    VAR_RE_FULL_PATTERN = VAR_RE_PATTERN + r'\s+0x[0-9a-fA-F]+\s+0x[0-9a-fA-F]+\s+\S+'
    VAR_RE_FULL_PATTERN_COMPILED = re.compile(VAR_RE_FULL_PATTERN)
    found_variable_in_map = False
    for index in range(0, len(self.map_lines)):
      line = self.map_lines[index]
      if VAR_RE_PATTERN_COMPILED.match(line):
        found_variable_in_map = True
        if not VAR_RE_FULL_PATTERN_COMPILED.match(line):
          line = line + self.map_lines[index + 1]
        break

    if found_variable_in_map:
      split_line_list = re.split(r' +', line.strip())
      self.variable_addr = split_line_list[1]
      self.variable_size = split_line_list[2]
      if (self.variable_addr[0:2].lower() != '0x') or (self.variable_size[0:2].lower() != '0x'):
        found_variable_in_map = False

    if not found_variable_in_map:
      self.LOG(LEVEL_ERROR, 'Cannot find the variable you want, Please check.')
    self.LOG(LEVEL_INFO, 'var:[{}], addr:[{}], size:[{}]'.format(self.variable_str, self.variable_addr, self.variable_size))


  def FindStartAddr(self):
    SECTION_RE_PATTERN = r'^\.\w+'
    SECTION_RE_PATTERN_COMPILED = re.compile(SECTION_RE_PATTERN)
    SECTION_RE_FULL_PATTERN = SECTION_RE_PATTERN + r'\s+0x[0-9a-fA-F]+\s+0x[0-9a-fA-F]'
    SECTION_RE_FULL_PATTERN_COMPILED = re.compile(SECTION_RE_FULL_PATTERN)
    found_fist_section = False
    for index in range(0, len(self.map_lines)):
      line = self.map_lines[index]
      line_next = self.map_lines[index + 1]
      if SECTION_RE_PATTERN_COMPILED.match(line):
        found_fist_section = True
        if not SECTION_RE_FULL_PATTERN_COMPILED.match(line):
          line = line + line_next
        break

    if found_fist_section:
      split_line_list = re.split(r' +', line.strip())
      self.bin_start_addr = split_line_list[1]
      if (self.bin_start_addr[0:2].lower() != '0x'):
        found_fist_section = False

    if not found_fist_section:
      self.LOG(LEVEL_ERROR, 'Cannot find the app bin start address in map file, Please check.')
    self.LOG(LEVEL_INFO, 'section:[{}], addr:[{}]'.format(split_line_list[0], self.bin_start_addr))
    pass


  def FindVariableBytes(self):
    self.variable_bytes_start = int(self.variable_addr, base=16) - int(self.bin_start_addr, base=16)
    self.variable_bytes_stop = self.variable_bytes_start + int(self.variable_size, base=16)
    self.variable_bytes = self.bin_bytes[self.variable_bytes_start: self.variable_bytes_stop]
    self.LOG(LEVEL_INFO, 'variable index in bin file: start [{}], stop [{}]'.format(self.variable_bytes_start, self.variable_bytes_stop))
    pass


  def GetVariableStruct(self):
    '''
      typedef struct {
          IIR_TYPE_T  type;
          float       gain;
          float       fc;
          float       Q;
      } IIR_PARAM_T;

      typedef struct {
          float   gain0;
          float   gain1;
          int     num;
          IIR_PARAM_T param[IIR_PARAM_NUM];
      } IIR_CFG_T;
    '''
    IIR_PARAM_NUM = 20
    UNPACK_FORMAT = '<ffi' + 'ifff' * IIR_PARAM_NUM
    unpack_data = struct.unpack(UNPACK_FORMAT, self.variable_bytes)
    eq = []
    offset = 0
    eq.append({'gain0': unpack_data[0]})
    eq.append({'gain1': unpack_data[1]})
    eq.append({'num': unpack_data[2]})
    offset = 3
    param = []
    for index in range(0, IIR_PARAM_NUM):
      type = unpack_data[offset]
      offset += 1
      gain = round(unpack_data[offset], 1)
      offset += 1
      fc = round(unpack_data[offset], 1)
      offset += 1
      Q = round(unpack_data[offset], 1)
      offset += 1
      param.append({index: [type, gain, fc, Q]})
    eq.append({'param': param})
    json_data = json.dumps(eq, sort_keys=False, indent=4)
    output_path = os.getcwd() + '/' + self.variable_str + '.txt'
    try:
      fd = open(output_path, 'w', encoding='utf8')
      fd.write(json_data)
      fd.close()
    except Exception as e:
      self.LOG(LEVEL_ERROR, 'Write variable to file failed: [{}], error=[{}]'.format(self.output_path, str(e)))
    self.LOG(LEVEL_INFO, 'Get variable struct successfull at [{}]'.format(output_path))
    pass


  def SetVariableStruct(self):
    IIR_PARAM_NUM = 20
    json_data = json.loads(self.variable_new_data_json)
    variable_new_bytes = b''
    gain0_bytes = struct.pack('<f', json_data[0].get('gain0'))
    gain1_bytes = struct.pack('<f', json_data[1].get('gain1'))
    num_bytes = struct.pack('<i', json_data[2].get('num'))
    params = json_data[3].get('param')
    params_bytes = gain0_bytes + gain1_bytes + num_bytes
    for index in range(0, IIR_PARAM_NUM):
      param = params[index]
      param_data = param.get(str(index))
      params_bytes += struct.pack('<ifff', param_data[0], param_data[1], param_data[2], param_data[3])

    app_bin_bytes_new = self.bin_bytes
    app_bin_bytes_new = self.bin_bytes[:self.variable_bytes_start] + params_bytes + self.bin_bytes[self.variable_bytes_stop:]
    new_bin_path = self.bin_path.rstrip('.bin') + '_new.bin'
    try:
      fd = open(new_bin_path, 'wb')
      fd.write(app_bin_bytes_new)
      fd.close()
    except Exception as e:
      self.LOG(LEVEL_ERROR, 'Can not open file: [{}], error=[{}]'.format(new_bin_path, str(e)))
    self.LOG(LEVEL_INFO, 'Write new bin file successfull at [{}]'.format(new_bin_path))
    pass


  def StartToOperate(self):
    self.ReadFiles()
    self.FindVariableAddr()
    self.FindStartAddr()
    self.FindVariableBytes()

    if self.workmode == 'r':
      self.GetVariableStruct()
    if self.workmode == 'w':
      self.SetVariableStruct()


if __name__ == '__main__':
  _log = _Log()
  LOG=_log.LOG
  LOG(LEVEL_INFO, "*********START*********")
  parser = argparse.ArgumentParser(description="Analyze the map")
  parser.add_argument("--bin", "-b", required=True, help="The app bin file you want to edit", metavar="APP_BIN_FILE")
  parser.add_argument("--map", "-m", required=True, help="The map file corresponding to the APP.bin you specified", metavar="MAP_FILE")
  parser.add_argument("--var", '-v', required=True, help="The variable you want to read or write.", metavar="VARIABLE")
  parser.add_argument("--mode", "-d", default='r', help="[r]: read the variable you specified [w]: write the variable based on the json file")
  parser.add_argument("--file", "-f", help="The json file you want you write to the app bin file")
  parser.add_argument("--out", "-o", help="The PATH you want to output the result, app_edited.bin", metavar="OUTPUT_FILE")
  args = parser.parse_args()

  if args.mode == 'w':
    if not args.file:
      LOG(LEVEL_ERROR, "Please specify the json file with '--input', whitch you want you write")
  editor = BinEditor(args.bin, args.map, args.var, args.mode, args.file, LOG)
  editor.StartToOperate()
  LOG(LEVEL_INFO, "*********END*********")
